#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
from typing import Dict, List, Tuple
import logging
from collections import defaultdict
from datetime import datetime, timedelta

class BehaviorAnalyzer:
    """VPN流量行为分析器"""
    
    def __init__(self, window_size: int = 100):
        self.logger = logging.getLogger(__name__)
        self.window_size = window_size
        self.packet_buffer = []
        self.flow_stats = defaultdict(lambda: {
            'packet_count': 0,
            'byte_count': 0,
            'start_time': None,
            'end_time': None,
            'inter_arrival_times': [],
            'protocols': set()
        })
        
    def analyze_behavior(self, packet: Packet, features: Dict[str, float]) -> Dict[str, float]:
        """
        分析数据包行为特征
        """
        # 更新数据包缓冲区
        self.packet_buffer.append((packet, features))
        if len(self.packet_buffer) > self.window_size:
            self.packet_buffer.pop(0)
            
        # 提取流标识
        flow_id = self._get_flow_id(packet)
        
        # 更新流统计信息
        self._update_flow_stats(flow_id, packet, features)
        
        # 计算行为特征
        behavior_features = {}
        behavior_features.update(self._calculate_flow_features(flow_id))
        behavior_features.update(self._calculate_temporal_patterns())
        behavior_features.update(self._calculate_protocol_patterns())
        
        return behavior_features
    
    def _get_flow_id(self, packet: Packet) -> str:
        """生成流标识"""
        if packet.haslayer(IP):
            src_ip = packet[IP].src
            dst_ip = packet[IP].dst
            
            if packet.haslayer(TCP):
                src_port = packet[TCP].sport
                dst_port = packet[TCP].dport
                proto = 'TCP'
            elif packet.haslayer(UDP):
                src_port = packet[UDP].sport
                dst_port = packet[UDP].dport
                proto = 'UDP'
            else:
                src_port = 0
                dst_port = 0
                proto = 'OTHER'
                
            # 确保源IP和目的IP的顺序一致
            if src_ip > dst_ip:
                src_ip, dst_ip = dst_ip, src_ip
                src_port, dst_port = dst_port, src_port
                
            return f"{src_ip}:{src_port}-{dst_ip}:{dst_port}-{proto}"
        return "unknown"
    
    def _update_flow_stats(self, flow_id: str, packet: Packet, features: Dict[str, float]):
        """更新流统计信息"""
        stats = self.flow_stats[flow_id]
        
        # 更新包计数和字节计数
        stats['packet_count'] += 1
        stats['byte_count'] += len(packet)
        
        # 更新时间戳
        current_time = datetime.fromtimestamp(packet.time)
        if stats['start_time'] is None:
            stats['start_time'] = current_time
        stats['end_time'] = current_time
        
        # 更新到达时间间隔
        if len(stats['inter_arrival_times']) > 0:
            last_time = datetime.fromtimestamp(self.packet_buffer[-2][0].time)
            inter_arrival = (current_time - last_time).total_seconds()
            stats['inter_arrival_times'].append(inter_arrival)
            
        # 更新协议集合
        if packet.haslayer(TCP):
            stats['protocols'].add('TCP')
        elif packet.haslayer(UDP):
            stats['protocols'].add('UDP')
        if packet.haslayer(TLS):
            stats['protocols'].add('TLS')
        if packet.haslayer(GRE):
            stats['protocols'].add('GRE')
        if packet.haslayer(ESP):
            stats['protocols'].add('ESP')
            
    def _calculate_flow_features(self, flow_id: str) -> Dict[str, float]:
        """计算流特征"""
        features = {}
        stats = self.flow_stats[flow_id]
        
        if stats['packet_count'] > 0:
            # 计算流持续时间
            if stats['start_time'] and stats['end_time']:
                duration = (stats['end_time'] - stats['start_time']).total_seconds()
                features['flow_duration'] = duration
                
            # 计算包大小统计特征
            features['avg_packet_size'] = stats['byte_count'] / stats['packet_count']
            
            # 计算到达时间间隔统计特征
            if stats['inter_arrival_times']:
                features['avg_inter_arrival'] = np.mean(stats['inter_arrival_times'])
                features['std_inter_arrival'] = np.std(stats['inter_arrival_times'])
                
            # 计算协议多样性
            features['protocol_diversity'] = len(stats['protocols'])
            
        return features
    
    def _calculate_temporal_patterns(self) -> Dict[str, float]:
        """计算时间模式特征"""
        features = {}
        
        if len(self.packet_buffer) > 1:
            # 计算包大小变化率
            sizes = [len(p[0]) for p in self.packet_buffer]
            features['size_variation'] = np.std(sizes) / np.mean(sizes) if np.mean(sizes) > 0 else 0
            
            # 计算时间间隔模式
            times = [p[0].time for p in self.packet_buffer]
            intervals = np.diff(times)
            features['time_interval_mean'] = np.mean(intervals)
            features['time_interval_std'] = np.std(intervals)
            
        return features
    
    def _calculate_protocol_patterns(self) -> Dict[str, float]:
        """计算协议模式特征"""
        features = {}
        
        if self.packet_buffer:
            # 计算协议转换频率
            protocol_changes = 0
            last_protocol = None
            
            for packet, _ in self.packet_buffer:
                current_protocol = None
                if packet.haslayer(TCP):
                    current_protocol = 'TCP'
                elif packet.haslayer(UDP):
                    current_protocol = 'UDP'
                    
                if last_protocol and current_protocol and last_protocol != current_protocol:
                    protocol_changes += 1
                last_protocol = current_protocol
                
            features['protocol_change_rate'] = protocol_changes / len(self.packet_buffer)
            
        return features 